"""Code for accepting self-signed information cards

Depends on the Python OpenID lbirary
"""
import sha
import base64
try:
    import OpenSSL
except ImportError:
    OpenSSL = None

import datetime
import re
import os
import subprocess
import tempfile
import time
import urlparse

import logging

from xml.etree import ElementTree

class TokenError(ValueError):
    """The token is malformed in some way
    """

class TokenExpiredError(TokenError):
    """The token has expired
    """

class UnexpectedAttributeValue(TokenError):
    """An attribute of a tag does not have the expected value
    """
    def __init__(self, tag_name, attr_name, expected_value, actual_value):
        TokenError.__init__(
            self, tag_name, attr_name, expected_value, actual_value)

    def __str__(self):
        return ("The %s node's attribute value %s is wrong. "
                "Expected %r, got %r" % tuple(self))

class SignedTokenError(TokenError):
    """The signed portion of the token is malformed
    """

XENC_NS = "http://www.w3.org/2001/04/xmlenc#"
XENC_ELEMENT_TYPE = "http://www.w3.org/2001/04/xmlenc#Element"
XENC_ENC_ALGO = "http://www.w3.org/2001/04/xmlenc#aes256-cbc"
XENC_KEYINFO_ENC_ALGO = "http://www.w3.org/2001/04/xmlenc#rsa-oaep-mgf1p"

DSIG_NS = "http://www.w3.org/2000/09/xmldsig#"
DSIG_RSA_SHA1 = "http://www.w3.org/2000/09/xmldsig#rsa-sha1"
DSIG_ENVELOPED_SIG = "http://www.w3.org/2000/09/xmldsig#enveloped-signature"
DSIG_SHA1 = "http://www.w3.org/2000/09/xmldsig#sha1"

CANON_EXCLUSIVE = "http://www.w3.org/2001/10/xml-exc-c14n#"

WSSE_NS = "http://docs.oasis-open.org" \
          "/wss/2004/01/oasis-200401-wss-wssecurity-secext-1.0.xsd"
WSSE_KEYID_VALUE_TYPE = "http://docs.oasis-open.org" \
                        "/wss/oasis-wss-soap-message-security-1.1" \
                        "#ThumbprintSHA1"

SAML_ASSERTION_1_0_NS = "urn:oasis:names:tc:SAML:1.0:assertion"
SAML_ASSERTION_1_1_NS = "http://docs.oasis-open.org" \
                        "/wss/oasis-wss-saml-token-profile-1.1#SAMLV1.1"

XMLSOAP_SELF_ISSUED = "http://schemas.xmlsoap.org" \
                      "/ws/2005/05/identity/issuer/self"

XMLSOAP_CLAIMS_NS = 'http://schemas.xmlsoap.org/ws/2005/05/identity/claims'

def nsTag(ns_uri, tag_name):
    """Generate a namespaced tag name in ElementTree format"""
    return '{%s}%s' % (ns_uri, tag_name)

def xencTag(tag_name):
    """Generate a namespaced tag name for the XMLENC namespace"""
    return nsTag(XENC_NS, tag_name)

def dsigTag(tag_name):
    """Generate a namespaced tag for the XMLDSIG namespace"""
    return nsTag(DSIG_NS, tag_name)

def wsseTag(tag_name):
    """Generate a namespaced tag for the WS-Security namespace"""
    return nsTag(WSSE_NS, tag_name)

def saml10Tag(tag_name):
    """Generate a namespaced tag for the SAML 1.0 Assertion namespace"""
    return nsTag(SAML_ASSERTION_1_0_NS, tag_name)

def saml11Tag(tag_name):
    """Generate a namespaced tag for the SAML 1.1 Assertion namespace"""
    return nsTag(SAML_ASSERTION_1_1_NS, tag_name)

def expectAttr(node, attr_name, expected):
    """Check that the given node's specified attribute has the expected value

    @raises UnexpectedAttributeValue: When the attribute does not have
        the expected value

    @returns: nothing
    """
    value = node.get(attr_name)
    if value != expected:
        raise UnexpectedAttributeValue(node.tag, attr_name, expected, value)

def getAttr(node, attr_name):
    """Check that the given node's specified attribute exists

    @raises UnexpectedAttributeValue: When the attribute is not present

    @returns: nothing
    """
    try:
        return node.attrib[attr_name]
    except KeyError:
        raise UnexpectedAttributeValue(node.tag, attr_name, None, None)

def findOrRaise(node, tag):
    """Utility function to find a single node given an XPath query
    (usually just a fully-qualifed tag name)

    @raises TokenError: Unless exactly one node matches the expression
    """
    nodes = node.findall(tag)
    if not nodes:
        raise TokenError('Missing %r tag' % (tag,))
    elif len(nodes) > 1:
        raise TokenError(' one %r tag found' % (tag,))

    return nodes[0]

def checkEncryptedToken(tree):
    """Check to see that the document that we've been passed is an
    encrypted document of the form that we expect for a Cardspace
    submission.
    """
    root_node = tree.getroot()
    if root_node.tag != xencTag('EncryptedData'):
        raise TokenError('The root is not an xenc:EncryptedData tag')

    expectAttr(root_node, 'Type', XENC_ELEMENT_TYPE)

    method_node = findOrRaise(root_node, xencTag('EncryptionMethod'))
    expectAttr(method_node, 'Algorithm', XENC_ENC_ALGO)

    keyinfo_node = findOrRaise(root_node, dsigTag('KeyInfo'))
    enckey_node = findOrRaise(keyinfo_node, xencTag('EncryptedKey'))
    enckey_method_node = findOrRaise(enckey_node, xencTag('EncryptionMethod'))
    expectAttr(enckey_method_node, 'Algorithm', XENC_KEYINFO_ENC_ALGO)

    enckey_keyinfo_node = findOrRaise(enckey_node, dsigTag('KeyInfo'))
    tokref_node = findOrRaise(enckey_keyinfo_node,
                              wsseTag('SecurityTokenReference'))

    keyid_node = findOrRaise(tokref_node, wsseTag('KeyIdentifier'))
    expectAttr(keyid_node, 'ValueType', WSSE_KEYID_VALUE_TYPE)

    # XXX: At this point, I think we're okay in ensuring that xmlenc
    # will be doing what we expect it to do. This should be confirmed.

def parseDate(xml_date):
    """Parse an XML date value to a Unix timestamp"""
    # There is no way to tell strptime to parse fractional seconds. We
    # don't really care about fractional seconds for the assertion
    # lifetime, anyway. This replace just drops the fractional seconds
    # part if it's present.
    xml_date = re.sub(r'\.\d+Z\Z', 'Z', xml_date, 1)

    # XXX: I bet there's a better way to get from time format to
    # datetime object.
    return datetime.datetime(
        *time.strptime(xml_date, '%Y-%m-%dT%H:%M:%SZ')[:6])

def getSignedTokenInfo(tree):
    """Check to see that the document that we've been passed is an
    encrypted document of the form that we expect for a Cardspace
    submission.
    """
    root_node = tree.getroot()
    if root_node.tag == saml10Tag('Assertion'):
        samlTag = saml10Tag
    elif root_node.tag == saml11Tag('Assertion'):
        samlTag = saml11Tag
    else:
        raise TokenError('The root is not a saml(v1.1 or 1.0):Assertion tag')

    token_info = TokenInformation()

    expectAttr(root_node, 'MajorVersion', '1')
    expectAttr(root_node, 'MinorVersion', '1')

    token_info.assertion_id = getAttr(root_node, 'AssertionID')

    # Issue instant is required by SAML spec, but not necessary to use
    # (?) since we have the conditions.
    token_info.issue_instant = parseDate(getAttr(root_node, 'IssueInstant'))

    conds_node = findOrRaise(root_node, samlTag('Conditions'))

    token_info.not_before = parseDate(getAttr(conds_node, 'NotBefore'))
    token_info.not_on_or_after = parseDate(getAttr(conds_node, 'NotOnOrAfter'))
    token_info.audience = findOrRaise(
        conds_node, '%s/%s' % (samlTag('AudienceRestrictionCondition'),
                               samlTag('Audience'))).text

    token_info.issuer = getAttr(root_node, 'Issuer')

    sig_node = findOrRaise(root_node, dsigTag('Signature'))

    signed_info_node = findOrRaise(sig_node, dsigTag('SignedInfo'))

    # "Uses the RSA signature method identified by the algorithm
    # identifier <http://www.w3.org/2000/09/xmldsig#rsa-sha1>."
    sig_meth_node = findOrRaise(signed_info_node, dsigTag('SignatureMethod'))
    expectAttr(sig_meth_node, 'Algorithm', DSIG_RSA_SHA1)

    # "Uses the exclusive canonicalization method identified by the
    # algorithm identifier <http://www.w3.org/2001/10/xml-exc-c14n#>
    # for canonicalizing the token content as well as the signature
    # content."
    canon_meth_node = findOrRaise(signed_info_node,
                                  dsigTag('CanonicalizationMethod'))
    expectAttr(canon_meth_node, 'Algorithm', CANON_EXCLUSIVE)

    # Now we need to find the assertion ID in the signature element so
    # we can make sure that the signature is for the right node.
    sig_reference_node = findOrRaise(signed_info_node, dsigTag('Reference'))

    # There are examples with both a URN and a relative fragment, so
    # this code handles both. XXX: make sure this is the right thing
    # to do. (The XMLDSIG spec is not very specific about how to
    # dereference the URI. The question is whether this is the correct
    # way to dereference a URN. The example suggests so, but maybe the
    # example is just missing the fragment character.)
    ref_uri = getAttr(sig_reference_node, 'URI')
    (scheme, _, _, _, _, fragment) = urlparse.urlparse(ref_uri)
    if scheme and scheme.lower() != 'urn':
        expected_assertion_id = ref_uri
    elif ref_uri.startswith('#'):
        expected_assertion_id = fragment
    else:
        raise TokenError('Expected a fragment or a URN. Got %r' % (ref_uri,))

    if expected_assertion_id != token_info.assertion_id:
        raise TokenError('Signature does not refer to the root node. '
                         'Expected %r, got %r' %
                         ('#' + token_info.assertion_id,
                          expected_assertion_id))

    # "No other transforms [...] are used in the enveloped signature."
    transform_nodes = sig_reference_node.findall(
        '%s/%s' % (dsigTag('Transforms'), dsigTag('Transform')))

    all_transforms = set([getAttr(node, 'Algorithm')
                          for node in transform_nodes])
    expected_transforms = set([CANON_EXCLUSIVE, DSIG_ENVELOPED_SIG])
    if all_transforms != expected_transforms:
        raise TokenError('Signature did not use expected transforms. '
                         'Expected %r, got %r' %
                         (expected_transforms, all_transforms))

    # "Uses the SHA1 digest method identified by the algorithm
    # identifier <http://www.w3.org/2000/09/xmldsig#sha1> for
    # digesting the token content being signed."
    digest_meth_node = findOrRaise(sig_reference_node, dsigTag('DigestMethod'))
    expectAttr(digest_meth_node, 'Algorithm', DSIG_SHA1)

    # According to <http://download.microsoft.com/download/6/c/3/
    # 6c3c2ba2-e5f0-4fe3-be7f-c5dcb86af6de/
    # infocard-guide-beta2-published.pdf>, this node must be there:
    #
    #     The ds:KeyInfo element is always present in the signature
    #     carrying the signing RSA public key in the form of a
    #     ds:RSAKeyValue child element.
    rsa_modulus_node = findOrRaise(sig_node,
        '%s/%s/%s/%s' % (
        dsigTag('KeyInfo'),
        dsigTag('KeyValue'),
        dsigTag('RSAKeyValue'),
        dsigTag('Modulus')))

    # XXX: "The signature key used in the issued token MUST be a
    # 2048-bit asymmetric RSA key which identifies the issuer."
    # We aren't checking the key length.
    token_info.rsa_modulus_b64 = rsa_modulus_node.text

    # Extract the attributes
    token_info.attributes = []
    attrs_node = findOrRaise(root_node, samlTag('AttributeStatement'))
    for attr_node in attrs_node.findall(samlTag('Attribute')):
        attr_ns = getAttr(attr_node, 'AttributeNamespace')
        attr_name = getAttr(attr_node, 'AttributeName')
        attr_val_node = findOrRaise(attr_node, samlTag('AttributeValue')).text
        token_info.attributes.append((attr_ns, attr_name, attr_val_node))



    # XXX: we do not implement this, as it's not necessary for what
    # we're doing now. Kim Cameron's example PHP code doesn't, either:
    #
    # When the subject confirmation method is "holder of key", the
    # subject confirmation key (also referred to as the proof key) MUST
    # be included in the token in the ds:KeyInfo child element under the
    # saml:SubjectConfirmation element. The proof key MUST be encoded in
    # the token as follows:
    #
    # * For symmetric key tokens, the proof key is encrypted to the
    #   recipient of the token in the form of a xenc:EncryptedKey
    #   child element. The default size of the key is 256 bits, but
    #   a different size may be specified by the relying party.
    #
    # * For asymmetric key tokens, the proof key is a public RSA
    #   key value specified as a ds:RSAKeyValue child element under
    #   ds:KeyValue element. The default size of the key is 2048
    #   bits.


    return token_info

class TokenInformation(object):
    """A container for the information parsed out of a self-signed
    infocard token"""

    def __init__(self):
        self.assertion_id = None
        self.issuer = None
        self.audience = None

        self.issue_instant = None
        self.not_before = None
        self.not_on_or_after = None

        self.rsa_modulus_b64 = None
        self.attributes = []

    def isExpired(self, now=None):
        """Is the current date within the window of this assertion"""
        if now is None:
            now = datetime.datetime.utcnow()

        slop = datetime.timedelta(minutes=5)
        expired =  (now >= (self.not_on_or_after + slop) or now < (self.not_before - slop))
        if expired:
            logging.error(
                "Got expired token: now=%s, not_before=%s, not_after=%s",
                now, self.not_before, self.not_on_or_after)
        return expired

    def getAttributes(self, attr_namespace=None):
        """Get attributes in this assertion, possibly for a specific
        namespace

        @param attr_namespace: Which namespace we want claims from
        """
        attrs = {}
        for (attr_ns, attr_name, attr_value) in self.attributes:
            if attr_namespace is None or attr_namespace == attr_ns:
                attrs[attr_name] = attr_value
        return attrs

    def checkAudience(self, cert_subject):
        """Check to make sure that the audience of this claim matches
        the specified cert subject's common name."""
        def err(text):
            raise ValueError('Unexpected audience: %s does not match %s: %s'
                             % (self.audience, cert_subject, text))

        audience_proto, audience_host = urlparse.urlparse(self.audience)[:2]
        # strip off the port number, if any
        audience_host = audience_host.rsplit(':', 1)[0]
        if audience_proto != 'https':
            err('Specified audience is not SSL')

        if cert_subject.startswith('*.'):
            if not audience_host.endswith(cert_subject[1:]):
                err('%r not under domain %r' % (audience_host, cert_subject,))
        elif audience_host != cert_subject:
            err('%r not under domain %r' % (audience_host, cert_subject,))

    def getRSAModulusSHA1(self):
        rsa_modulus = base64.b64decode(self.rsa_modulus_b64)
        return sha.new(rsa_modulus).digest()

    def getStandardClaim(self, claim_name):
        for (attr_ns, attr_name, attr_value) in self.attributes:
            if (attr_ns == XMLSOAP_CLAIMS_NS and
                attr_name == claim_name):
                return attr_value

        return None

class TokenProcessor(object):
    xmlsec_filename = 'xmlsec1'

    def __init__(self, cert_filename, temp_dir=None, common_name=None):
        self.cert_filename = cert_filename
        self.x509_cert = self.getX509()
        self.temp_dir = temp_dir
        self.common_name = common_name

    def getX509(self):
        """Parse the specified certificate file"""
        if OpenSSL is None:
            return None

        cert_data = file(self.cert_filename).read()
        return OpenSSL.crypto.load_certificate(
            OpenSSL.crypto.FILETYPE_PEM, cert_data)

    def getCommonName(self):
        if not self.common_name and self.x509_cert:
            self.common_name = self.x509_cert.get_subject().commonName
        return self.common_name

    def processXML(self, xml_text):
        (outfd, outname) = tempfile.mkstemp(dir=self.temp_dir)
        outfile = os.fdopen(outfd, 'w+')
        outfile.write(xml_text)
        outfile.close()
        return self.process(outname)

    def process(self, encrypted_token_filename):
        # First check to see that the encrypted document has the
        # expected form
        tree = ElementTree.parse(file(encrypted_token_filename))
        checkEncryptedToken(tree)

        signed_token_filename = self.decrypt(encrypted_token_filename)

        tree = ElementTree.parse(file(signed_token_filename))
        token_info = getSignedTokenInfo(tree)

        # Call xmlsec1 to verify the signature
        self.checkSig(signed_token_filename)

        token_info.checkAudience(self.common_name)

        if token_info.isExpired():
            raise TokenExpiredError()

        return token_info

    def decrypt(self, encrypted_token_filename):
        (outfd, outname) = tempfile.mkstemp(dir=self.temp_dir)
        outfile = os.fdopen(outfd, 'w+')

        # Then invoke xmlsec on the encrypted document to decrypt
        proc = subprocess.Popen(
            [self.xmlsec_filename,
             '--decrypt',
             '--privkey-pem', self.cert_filename,
             encrypted_token_filename],
            stdout=outfile,
            stderr=subprocess.PIPE,
            )
        (_, err_out) = proc.communicate()
        status = proc.wait()
        if status != 0:
            raise TokenError('xmlsec1 failed: ' + err_out)

        outfile.flush()
        outfile.close()

        return outname

    def checkSig(self, signed_token_filename):
        for saml_ns in [SAML_ASSERTION_1_0_NS, SAML_ASSERTION_1_1_NS]:
            proc = subprocess.Popen(
                [self.xmlsec_filename,
                 '--verify',

                 # Since we don't have the DTD or schema for the SAML
                 # assertion, we need to tell it what the id attribute
                 # is called
                 '--id-attr:AssertionID', saml_ns + ':Assertion',

                 # Operate on the root document
                 '--node-xpath', '/',

                 # Only allow reference to the same document
                 '--enabled-reference-uris', 'same-doc',

                 signed_token_filename],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                )

            (_, err_out) = proc.communicate()
            status = proc.wait()
            if status != 0:
                continue

            expected_output = """OK
SignedInfo References (ok/all): 1/1
Manifests References (ok/all): 0/0
"""
            if err_out != expected_output:
                continue

            return
        else:
            # err_out is set in each error case
            raise TokenError('xmlsec1 failed: ' + err_out)
